#!/usr/bin/env python

import collections
import cPickle
import math
import optparse
import os
import sys

import minuit
import numpy
from matplotlib import pyplot
import scipy
from scipy import integrate

import charmicat
from physics.function import breit_wigner

from iiipi.pwa.guru import PathGuru
from iiipi.pwa import mdfit
from iiipi import textable

def calculate_yield(x, scale, mass, width, l):
    amp = breit_wigner(x, 775.5, 139.57, 
                       mass, width, l)

    return scale*(amp*amp.conjugate()).real

def calculate_dphi(x, offset, m1, w1, m2, w2, l1=1, l2=2):
    amp1 = breit_wigner(x, 775.5, 139.57, 
                       m1, w1, l1)

    amp2 = breit_wigner(x, 775.5, 139.57, 
                        m2, w2, l2)

    return math.atan2(amp1.imag, amp1.real) - math.atan2(amp2.imag, amp2.real) + offset


def chi2(scale1, m1, gamma1, scale2, m2, gamma2, offset):
    chi2 = 0
    global x
    global emat
    global measured_vals 
    global bins_in_range

    for i, bin in enumerate(bins_in_range):
        calc_y1 = calculate_yield(x[i], scale1, m1, gamma1, 1)
        calc_y2 = calculate_yield(x[i], scale2, m2, gamma2, 2)
        calc_dphi = calculate_dphi(x[i], offset, m1, gamma1, m2, gamma2)

        diff = numpy.array([measured_vals['y1'][i] - calc_y1,
                            measured_vals['y2'][i] - calc_y2,
                            measured_vals['dphi'][i] - calc_dphi])

        chi2 += numpy.dot(diff.T, numpy.dot(emat[bin], diff))


    return chi2



def dphi_chi2(m1, gamma1, m2, gamma2, offset):
    global x
    global emat
    global measured_vals 
    global bins_in_range

    chi2 = 0
    
    for i, bin in enumerate(bins_in_range):
        calc_dphi = calculate_dphi(x[i], offset, m1, gamma1, m2, gamma2)

        chi2 += (measured_vals['dphi'][i] - calc_dphi)**2 * emat[bin][2][2]

    return chi2



def main():
    op = optparse.OptionParser()

    op.add_option('-i', '--init-params', dest='initpars', default=None, 
                  help='Specify initial parameters in FILE.py', metavar='FILE.py')
    op.add_option('-r', '--range', dest='range', nargs=2, type=int, default=(1000, 2000),
                  help='Specify the fit range LOW HIGH', metavar='LOW HIGH')


    opts, args = op.parse_args()

    if len(args) < 4:
        msg = "Four command-line arguments are required:\n"
        msg += "\t- a text file with corrected yields for some wave;\n"
        msg += "\t- a text file with corrected yields for some other wave;\n"
        msg += "\t- a text file with the offset diffs between the two waves;\n"
        msg += "\t- the path to the fit directory where these results come from.\n"
        raise charmicat.BadCommandLineError(msg, -1)

    y1path = args.pop(0)
    y2path = args.pop(0)
    dphi_path = args.pop(0)
    fitdir = args.pop(0)

    if not os.path.isdir(fitdir):
        msg = "The last argument must point to a directory."
        raise charmicat.BadCommandLineError(msg, -2)

    guru = PathGuru(fitdir)

    bng = guru.fit_binning


    global bins_in_range
    bins_in_range = ["{0}{1}-{2}".format(bng.tag, bin, bin+bng.width) 
                     for bin in xrange(opts.range[0], opts.range[1], bng.width)] 

#    bins_in_range = '''Xbin1000-1020 Xbin1020-1040 Xbin1040-1060 Xbin1060-1080 
#                       Xbin1080-1100 Xbin1100-1120 Xbin1120-1140 Xbin1140-1160 
#                       Xbin1160-1180 Xbin1180-1200 Xbin1200-1220 Xbin1220-1240 
#                       Xbin1240-1260 Xbin1260-1280 Xbin1280-1300 Xbin1300-1320  
#                       Xbin1320-1340 Xbin1340-1360 Xbin1360-1380 Xbin1380-1400'''.split()

    indices = [bng.dirs.index(bin) for bin in bins_in_range]


    xfull, dphi_full, dphi_err_full = numpy.loadtxt(dphi_path).T

    x1, y1full, y1err_full = numpy.loadtxt(y1path, usecols=(1,2,3)).T
    x2, y2full, y2err_full = numpy.loadtxt(y2path, usecols=(1,2,3)).T

    global x
    x = xfull[indices]
    dphi_err = dphi_err_full[indices]
    y1err = y1err_full[indices]
    y2err = y2err_full[indices]


    # measured values (for each bin):
    # y1 y2 dphi 

    global measured_vals
    measured_vals = {}
    measured_vals['y1'] = y1full[indices]
    measured_vals['y2'] = y2full[indices]
    measured_vals['dphi'] = dphi_full[indices]


    print guru.fitdir

    global emat
    emat = mdfit.transform_errors(y1path, y2path, guru, bins=bins_in_range)


    # parameters:
    # scale1 m1 gamma1 scale2 m2 gamma2


    m = minuit.Minuit(chi2)
    m.maxcalls = 50000
    #m = minuit.Minuit(dphi_chi2)

    if opts.initpars:
        with open(opts.initpars) as fp:
            init_params = eval(fp.read())
        
        m.values = init_params

    else:
        m.values['scale1'] = 1.31e5
        m.values['scale2'] = 7e5
        m.values['m1']     = 1190
        m.values['m2']     = 1310
        m.values['gamma1'] = 303
        m.values['gamma2'] = 100
        m.values['offset'] = 0.0


#    dphi_calc = [calculate_dphi(ix, m.values['offset'], m.values['m1'], m.values['gamma1'], m.values['m2'], m.values['gamma2']) for ix in x]
    
#    print measured_vals['dphi'] - dphi_calc
    
#    dphi_calc = [calculate_dphi(ix, 3.669, m.values['m1'], m.values['gamma1'], m.values['m2'], m.values['gamma2']) for ix in x]
#    print measured_vals['dphi'] - dphi_calc

#    return

#    m.limits['offset'] = (0.75, 2)

    initial_chi2 = chi2(**m.values)
    #print dphi_chi2(**m.values)

    #m.scan()
    #m.printMode = 1
    m.migrad()

    print 'offset', m.values['offset']

    for k in m.values:
        m.fixed[k] = True

    m.fixed['offset'] = False

    m.migrad()
    

    final_chi2 = chi2(**m.values)
    #print dphi_chi2(**m.values)

    parnames = sorted(m.values.keys())

    for pn in parnames:
        print pn, m.values[pn]

    print '{0:.3f} {1:.3f}'.format(initial_chi2, final_chi2)

#    xcurve = numpy.linspace(1000, 2000, 10000)
    xcurve = numpy.linspace(opts.range[0], opts.range[1], 10000)

    figx = 8
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.8, 0.8))

    nsignal = integrate.quad(lambda x: calculate_yield(x, m.values['scale1'], m.values['m1'], m.values['gamma1'], 1), 
                             -scipy.inf, scipy.inf)

    info = [m.values['scale1'], m.errors['scale1'], m.values['m1'], m.errors['m1'], m.values['gamma1'], m.errors['gamma1'], 
            math.fabs(initial_chi2), math.fabs(final_chi2), 
            nsignal[0]]
    
    fptex = textable.write_table_bw(info)


    pyplot.text(0.95, 0.95,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'alpha':'0.7', 'color':'#FFFFFF'})


    pyplot.errorbar(xfull, y1full, yerr=y1err_full, fmt='k.')
    pyplot.plot(xcurve, [calculate_yield(ix, m.values['scale1'], m.values['m1'], m.values['gamma1'], 1) for ix in xcurve],
                'r-', lw=2, alpha=0.75)

    pyplot.ylim(0, pyplot.ylim()[1])

    if 'm1p' in y1path:
        pyplot.title(r'Mass-Dependent Fit to $1^{++}1^+S$')
    elif 'm1m' in y1path:
        pyplot.title(r'Mass-Dependent Fit to $1^{++}1^-S$')

    pyplot.xlabel(r"M($\pi^+\pi^+\pi^-$) (MeV$/c^2$)")
    pyplot.ylabel(r"Events per 20 MeV$/c^2$")

    pyplot.savefig("a1intensity.png")
    pyplot.savefig("a1intensity.pdf")

    pyplot.clf()

    figx = 8
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.8, 0.8))

    nsignal = integrate.quad(lambda x: calculate_yield(x, m.values['scale2'], m.values['m2'], m.values['gamma2'], 1), 
                             -scipy.inf, scipy.inf)

    info = [m.values['scale2'], m.errors['scale2'], m.values['m2'], m.errors['m2'], m.values['gamma2'], m.errors['gamma2'], 
            math.fabs(initial_chi2), math.fabs(final_chi2), 
            nsignal[0]]
    
    fptex = textable.write_table_bw(info)

    pyplot.text(0.95, 0.95,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'alpha':'0.7', 'color':'#FFFFFF'})

    pyplot.errorbar(xfull, y2full, yerr=y2err_full, fmt='k.')
    pyplot.plot(xcurve, [calculate_yield(ix, m.values['scale2'], m.values['m2'], m.values['gamma2'], 2) for ix in xcurve],
                'r-', lw=2, alpha=0.75)

    pyplot.ylim(0, pyplot.ylim()[1])

    if 'm1p' in y1path:
        pyplot.title(r'Mass-Dependent Fit: $2^{++}1^+D$ Intensity')
    elif 'm1m' in y1path:
        pyplot.title(r'Mass-Dependent Fit to $2^{++}1^-D$')

    pyplot.xlabel(r"M($\pi^+\pi^+\pi^-$) (MeV$/c^2$)")
    pyplot.ylabel(r"Events per 20 MeV$/c^2$")

    pyplot.savefig("a2intensity.png")
    pyplot.savefig("a2intensity.pdf")

    pyplot.clf()

    figx = 8
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.8, 0.8))

    pyplot.errorbar(xfull,  dphi_full, yerr=dphi_err_full, fmt='k.')
    pyplot.plot(xcurve, [calculate_dphi(ix, m.values['offset'], m.values['m1'], m.values['gamma1'], 
                                        m.values['m2'], m.values['gamma2']) for ix in xcurve],
                'r-', lw=2, alpha=0.75)

    if 'm1p' in y1path:
        pyplot.title(r'Mass-Dependent Fit: $1^{++}1^+S - 2^{++}1^+D$ $\delta\phi$')
    elif 'm1m' in y1path:
        pyplot.title(r'Mass-Dependent Fit: $1^{++}1^-S - 2^{++}1^-D$ $\delta\phi$')

    pyplot.ylim(-math.pi, math.pi)

    pyplot.gca().set_yticks([-math.pi, -0.75*math.pi, -0.5*math.pi, -0.25*math.pi,
                              0., 0.25*math.pi, 0.5*math.pi, 0.75*math.pi, math.pi])
    pyplot.gca().set_yticklabels([r'-$\pi$', r'$-\frac{3}{4}\pi$', r'$-\frac{1}{2}\pi$', r'$-\frac{1}{4}\pi$', 
                              r'$0$', r'$\frac{1}{4}\pi$', r'$\frac{1}{2}\pi$', r'$\frac{3}{4}\pi$', r'$\pi$'])

    pyplot.xlabel(r"M($\pi^+\pi^+\pi^-$) (MeV$/c^2$)")
    pyplot.ylabel(r"$\delta\phi$ radians per 20 MeV$/c^2$")

    pyplot.savefig("a1a2_phase.png")
    pyplot.savefig("a1a2_phase.pdf")





if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0}' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
